import argparse
import os.path
import traceback

import cv2
import matplotlib.pyplot as plt
import numpy as np

import math
import time

from Common.ControlEnum import forcemode
from Common.CreateFolder import FolderManager
from Common.INIFile import IniConfig
from Hardware.Gripper2F85 import  Robotiq2F85
from Hardware.UrRobot import UrRobot
from Hardware import FT300
from Common import INIFile
from Hardware import FTur5
from Algorithm.AlgorithmModel import ControlResult
from Algorithm.M_my_test_MNR import my_test
from Hardware.realsense.realsense import Camera


def plot(data,path):
    # 桌面硬表面测试
    # ur5.moveJ([-258.8 * a2r, -131.33 * a2r, -63.38 * a2r, -50.85 * a2r, 85.1 * a2r, -41.07 * a2r], 0.8, 0.1)
    count = len(data)
    plt.figure(1)
    # plt.axis("equal")
    posx = []
    posy = []
    posz = []
    fz = []
    fx = []
    fy = []
    tz = []
    tx = []
    ty = []
    vel = []
    for i in range(count):
        posx.append(data[i]['actualTcpPose'][0])
        posy.append(data[i]['actualTcpPose'][1])
        posz.append(data[i]['actualTcpPose'][2])

        tx.append(data[i]['force'][3])
        ty.append(data[i]['force'][4])
        tz.append(data[i]['force'][5])
        fx.append(data[i]['force'][0])
        fy.append(data[i]['force'][1])
        fz.append(data[i]['force'][2])

        vel.append(np.linalg.norm(data[i]['actualTcpSpeed'][1]))
    # plt.plot(posx,posy)
    # np.savetxt("force.txt",fz)
    # np.savetxt("pos.txt",posz)
    axisx = range(1, len(vel) + 1)
    # plt.figure(2)
    plt.plot(axisx, fz, label="fz")
    plt.plot(axisx, fx, label="fx")
    plt.plot(axisx, fy, label="fy")
    plt.legend()
    if(path!=""):
        plt.savefig(os.path.join(path,"force"))
    plt.show()

    plt.figure(2)
    plt.plot(axisx, tz, label="tz")
    plt.plot(axisx, tx, label="tx")
    plt.plot(axisx, ty, label="ty")
    plt.legend()
    if(path!=""):
        plt.savefig(os.path.join(path,"torque"))
    plt.show()

def move_contact(ur5,goalforce):
    fz=ur5.GetForceBase()[2]
    while(fz<goalforce):
        pose=ur5.getActualTCPPose()
        pose[2]=pose[2]-0.0005
        ur5.servoL(pose)
        # print(pose)
        time.sleep(0.01)
        fz = ur5.GetForceBase()[2]
        # print(fz)
    print("当前接触力",ur5.GetForceBase()[2])

def save_pose_info(folder,grippose,initpose):
    from configparser import ConfigParser

    file=os.path.join(folder,"info.ini")
    config = ConfigParser()

    config.add_section('position')
    config.set('position', 'grippose', str(grippose))
    config.set('position', 'initpose', str(initpose))
    with open(file, 'w') as configfile:
        config.write(configfile)

def main(args):

    # region 获取配置文件
    a2r = math.pi / 180
    cf=args['configfile']
    config = INIFile.IniConfig(cf)
    robotip = config.read('hardware', 'robotip')
    useFt = config.read('hardware', 'UseForceSensor')
    robot_type = config.read('hardware', 'robot')
    teachMode = forcemode(args["teach_mode"])
    dt=0.008 if robot_type == "ur5" else 0.002
    print(robotip,useFt,dt,teachMode)
    isReplaySave=False
    if(teachMode==forcemode.ReplayMode and args["is_replay_record"]):
        isReplaySave=True
        if(os.path.exists(args["replay_folder"]+"//images")==False):
            os.makedirs(args["replay_folder"]+"//images")
    if (teachMode == forcemode.TeachMode):
        # 创建文件夹
        target_directory = args["demo_path_root"]  # Replace this with the desired target directory
        folder_manager = FolderManager(args["teach_folder"], target_directory)
        saveFolder = folder_manager.create_new_folder()
        args["save_folder"]=saveFolder#新建拖动示教保存目录
    # endregion

    # region 硬件初始化
    ur5 = UrRobot(robotip=robotip, forcesensor=FT300.RobotiqFt300() if useFt == "True" else FTur5.Ftur5(),dt=dt)
    ur5.init()

    gripper = Robotiq2F85(serial_number='COM3')
    gripper.init()
    gripper.go_to(85,150,20)

    camera = Camera()
    camera.Open()

    # robotip="192.168.199.128"
    # ur5 = UrRobot(robotip=robotip, forcesensor=FTur5.Ftur5(), dt=dt)

    # endregion

    #region 初始化移动
    if(teachMode==forcemode.TeachMode):
        gripperpose=np.array([-27.20,-109.24,-76.59,-84.56,91.58,65.67])
        initpose = np.array([-26.68, -109, -73 , -90 , 92, 73])
        pose1=(gripperpose*a2r).tolist()
        pose2=(initpose*a2r).tolist()
        # pose2=[-0.4743804906920588, -1.9135789918865829, -1.2868312574954193, -1.5482815794441696, 1.6074482410867774, 0.9054768159346582]
    if(teachMode==forcemode.ReplayMode):
        #去文件夹里面找文件获取
        import ast
        config=IniConfig(os.path.join(args["replay_folder"],"info.ini"))
        pose1=ast.literal_eval(config.read("position","grippose"))
        pose2 = ast.literal_eval(config.read("position", "initpose"))

    pose=ur5.getActualTCPPose()
    pose[2]+=0.02
    ur5.moveL(pose,0.1,0.1)
    ur5.moveJ(pose1, 0.1, 0.1)
    gripper.go_to(0,150,235)
    # time.sleep(2)
    pose = ur5.getActualTCPPose()
    pose[2] += 0.02
    ur5.moveL(pose, 0.1, 0.1)
    ur5.moveJ(pose2, 0.05, 0.1)

    ur5.FtZero()
    time.sleep(0.5)
    move_contact(ur5,25)
    #向下缓慢移动，直至接触

    #endregion

    # 3  核心部分
    initPose = ur5.getActualTCPPose()
    algorithm = my_test(initPose,args,dt,teachMode)
    ur5.setControlFun(algorithm.Move)
    ur5.setStartRealTimeControl(True)
    ur5.enableLog()
    count=0
    while (True):
        time.sleep(0.001)
        _,image=camera.GetImage()
        if(isReplaySave):
            cv2.imwrite(os.path.join(args["replay_folder"],"images",str(count)+".png"),image)
            # cv2.imshow("1",image)
            cv2.waitKey(30)
            count+=1
            algorithm.controller.add_muti_sensor_time()
        # print(123)
        if (ur5.getRealTimeControlResult() == ControlResult.Finish):
            break

    data = ur5.disableLog()
    plot(data,args["replay_folder"] if teachMode==forcemode.ReplayMode else "")
    if(teachMode==forcemode.TeachMode):
        save_pose_info(args["save_folder"],pose1,pose2)
    time.sleep(1)
    gripper.go_to(85,100,20)
    ur5.disconnect()


if __name__ == '__main__':
    #示教模式，将teach_mode改为1
    #复现模式，将teach_mode改为2，同时选择复现轨迹replay_folder，并选择是否记录is_replay_record
    parser = argparse.ArgumentParser()
    parser.add_argument('--teach_mode', type=int, default="2", required=False)#TeachMode=1,ReplayMode=2,FreeMode=3
    parser.add_argument('--replay_folder', type=str, default="demo13", required=False)#复现模式的轨迹文件

    parser.add_argument('--configfile',type=str,default="config.ini", required=False)
    parser.add_argument('--demo_path_root', type=str, default="result/demotraj", required=False)
    parser.add_argument('--teach_folder', type=str, default="demo", required=False)#示教模式保存的文件夹名字
    parser.add_argument('--is_replay_record', type=bool, default=True, required=False)#复现模式是否记录数据
    args=vars(parser.parse_args())#变为字典
    args["replay_folder"]=os.path.join(args["demo_path_root"],args["replay_folder"])
    main(args)
